import argparse

"""
Here are the param for the training

"""


def get_args():
    parser = argparse.ArgumentParser("Reinforcement Learning experiments for multiagent environments")
    # Environment
    parser.add_argument("--scenario-name", type=str, default="uav_test", help="name of the scenario script")
    parser.add_argument("--max-episode-len", type=int, default=100, help="maximum episode length")
    parser.add_argument("--time-steps", type=int, default=2000000, help="number of time steps")
    parser.add_argument("--num_uavs", type=int, default=3, help="number of agents")
    parser.add_argument("--grid_size", type=int, default=10, help="the size of the grid")
    parser.add_argument("--num_position_action", type=int, default=4, help="number of position_action")
    parser.add_argument("--num_states",type=int, default=2, help="number of states")
    parser.add_argument("--num_other_uav_states",type=int, default=8, help="number of the other uav states")
    parser.add_argument("--num_actions",type=int, default=4, help="number of actions")
    parser.add_argument("--num_rewards",type=int, default=1, help="number of rewards")
    # The IB model
    parser.add_argument("--max_latent_dim", type=int, default=8, help="max latent dim")
    parser.add_argument("--min_latent_dim", type=int, default=4, help="min latent dim")
    parser.add_argument("--initial_ber", type=float, default=0.001, help="initial ber")
    parser.add_argument("--max_ber", type=float, default=0.01, help="max ber")
    parser.add_argument("--alpha", type=float, default=0.1, help="alpha")
    parser.add_argument("--beta", type=float, default=0.5, help="beta")
    parser.add_argument("--lr", type=float, default=1e-4, help="learning rate")
    parser.add_argument("--low_ber_threshold", type=float, default=0.005, help="low ber threshold")
    parser.add_argument("--high_ber_threshold", type=float, default=0.005, help="high ber threshold")   
    parser.add_argument("--attention_heads", type=int, default=4, help="attention heads")
    parser.add_argument("--feature_dims", type=list, default=[512, 256, 128], help="feature dims")
    # Core training parameters
    parser.add_argument("--lr-actor", type=float, default=1e-4, help="learning rate of actor")
    parser.add_argument("--lr-critic", type=float, default=1e-3, help="learning rate of critic")
    parser.add_argument('--lr_step_size', type=int, default=100, help='Number of training steps after which the learning rate decreases')
    parser.add_argument('--lr_gamma', type=float, default=0.99, help='Factor by which the learning rate will be reduced')
    parser.add_argument("--epsilon", type=float, default=0.01, help="epsilon greedy")
    parser.add_argument("--noise_rate", type=float, default=0.1, help="noise rate for sampling from a standard normal distribution ")
    parser.add_argument("--gamma", type=float, default=0.95, help="discount factor")
    parser.add_argument("--tau", type=float, default=0.005, help="parameter for updating the target network")
    parser.add_argument("--l2-critic", type=float, default=1e-3, help="L2 regularization for critic")
    parser.add_argument("--buffer-size", type=int, default=int(5e5), help="number of transitions can be stored in buffer")
    parser.add_argument("--batch-size", type=int, default=64, help="number of episodes to optimize at the same time")
    # Checkpointing
    parser.add_argument("--save-dir", type=str, default="./model", help="directory in which training state and model should be saved")
    parser.add_argument("--save-rate", type=int, default=200, help="save model once every time this many episodes are completed")
    parser.add_argument("--model-dir", type=str, default="", help="directory in which training state and model are loaded")
    parser.add_argument("--history_length_input",type=int,default=5,help="input_nums")
    # Evaluate
    parser.add_argument("--evaluate-episodes", type=int, default=50, help="number of episodes for evaluating")
    parser.add_argument("--evaluate-episode-len", type=int, default=100, help="length of episodes for evaluating")
    parser.add_argument("--evaluate", type=bool, default=False, help="whether to evaluate the model")
    parser.add_argument("--uavtrain", type=bool, default=True, help="whether to train the uav model")
    parser.add_argument("--evaluate-rate", type=int, default=1000, help="how often to evaluate model")
    args = parser.parse_args()


    # add the para
    parser.add_argument('--critic_lr',type=float,default=0.001)
    parser.add_argument('--res_connect', type=bool, default=True, help='Enable residual connection')
    return args
